{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPLt3b5Obzy9cBPO2FGOJMN",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Chat-Haruhi-Suzumiya/blob/main/notebook/Erotic%E7%9A%84%E5%B0%8F%E8%AF%B4%E6%95%B0%E6%8D%AE%E7%94%9F%E6%88%90.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "- [x] 统计换行数量\n",
        "- [x] 统计引号数量\n",
        "- [ ] prompt改写成多样化的prompt"
      ],
      "metadata": {
        "id": "S024bFjfrRs1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0Z_hevNCrRIJ",
        "outputId": "667b58b7-9423-4ec1-b32d-1cfc0656ddb0"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8_c1yqgApHy3"
      },
      "outputs": [],
      "source": [
        "!ls /content/drive/MyDrive/Wuxia/erotica"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q tiktoken\n",
        "import tiktoken\n",
        "\n",
        "enc = tiktoken.get_encoding(\"cl100k_base\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0Aic_8sw_QPQ",
        "outputId": "a297d36c-a33f-4b6d-a979-6a949d5e3f80"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/2.0 MB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.1/2.0 MB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.9/2.0 MB\u001b[0m \u001b[31m12.9 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m20.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "llmx 0.0.15a0 requires cohere, which is not installed.\n",
            "llmx 0.0.15a0 requires openai, which is not installed.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 定义strong_divide\n",
        "def divide_str(s, sep=['\\n', '.', '。'], consider_quote = True ):\n",
        "    mid_len = len(s) // 2  # 中心点位置\n",
        "    best_sep_pos = len(s) + 1  # 最接近中心点的分隔符位置\n",
        "    best_sep = None  # 最接近中心点的分隔符\n",
        "    for curr_sep in sep:\n",
        "        sep_pos = s.rfind(curr_sep, 0, mid_len + mid_len - 2)  # 从中心点往左找分隔符\n",
        "        if sep_pos > 0 and abs(sep_pos - mid_len) < abs(best_sep_pos -\n",
        "                                                        mid_len):\n",
        "            if consider_quote == True:\n",
        "                if s[sep_pos + 1] in ['\"', \"'\",\"“\",\"」\"]:\n",
        "                    sep_pos += 1\n",
        "\n",
        "            best_sep_pos = sep_pos\n",
        "            best_sep = curr_sep\n",
        "    if not best_sep:  # 没有找到分隔符\n",
        "        return s, ''\n",
        "    return s[:best_sep_pos + 1], s[best_sep_pos + 1:]\n",
        "\n",
        "\n",
        "def strong_divide(s, balance_ratio = 0.25):\n",
        "    left_bar = int(len(s) * balance_ratio)\n",
        "    left, right = divide_str(s, '\\n')\n",
        "    if len(left) >= left_bar and len(right) >= left_bar:\n",
        "        return left, right\n",
        "\n",
        "    whole_sep = ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                 '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                 '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                 '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "               '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                 '·', '…']\n",
        "    left, right = divide_str(s, sep=whole_sep)\n",
        "    if len(left) >= left_bar and len(right) >= left_bar:\n",
        "        return left, right\n",
        "\n",
        "    mid_len = len(s) // 2\n",
        "    return s[:mid_len], s[mid_len:]\n",
        "\n",
        "def super_divide( s, max_len_token = 300, iter_deep = 1):\n",
        "    n_token = len(enc.encode(s))\n",
        "    if n_token <= max_len_token or iter_deep > 2:\n",
        "        return [s]\n",
        "\n",
        "    left, right = strong_divide(s)\n",
        "    return super_divide(left, max_len_token, iter_deep+1) + \\\n",
        "            super_divide(right, max_len_token, iter_deep+1)"
      ],
      "metadata": {
        "id": "GLK23YfQ_Sbu"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "对于 /content/drive/MyDrive/Wuxia/erotica 目录下的所有文件夹\n",
        "\n",
        "如 /content/drive/MyDrive/Wuxia/erotica/例子小说1\n",
        "\n",
        "这个时候我需要定义\n",
        "\n",
        "source = \"例子小说1\"\n",
        "\n",
        "对于这个目录下面，有很多txt\n",
        "\n",
        "我希望对每个txt文件进行读取\n",
        "\n",
        "保存到一个list of dict , datas中，保存source 和text两个字段\n",
        "\n",
        "请用python为我实现"
      ],
      "metadata": {
        "id": "HI6F6tdP3oHE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "# 定义目录路径\n",
        "directory_path = \"/content/drive/MyDrive/Wuxia/erotica\"\n",
        "\n",
        "# 初始化用于保存数据的列表\n",
        "datas = []\n",
        "\n",
        "from tqdm import tqdm\n",
        "\n",
        "count = 0\n",
        "\n",
        "import time\n",
        "\n",
        "# 遍历目录下的所有文件夹\n",
        "for folder in tqdm(os.listdir(directory_path)):\n",
        "    folder_path = os.path.join(directory_path, folder)\n",
        "\n",
        "    count += 1\n",
        "    # if count > 10:\n",
        "    #     break\n",
        "\n",
        "    # 检查当前路径是否为文件夹\n",
        "    if os.path.isdir(folder_path):\n",
        "\n",
        "        # 遍历文件夹下的所有txt文件\n",
        "        for filename in os.listdir(folder_path):\n",
        "            if filename.endswith(\".txt\"):\n",
        "                file_path = os.path.join(folder_path, filename)\n",
        "\n",
        "                # 读取txt文件内容\n",
        "                with open(file_path, \"r\", encoding=\"utf-8\") as file:\n",
        "                    text_content = file.read()\n",
        "\n",
        "                # 将source和text添加到datas列表中\n",
        "                datas.append({\"source\": folder, \"text\": text_content})\n",
        "\n",
        "                time.sleep(0.05)"
      ],
      "metadata": {
        "id": "aOGYrFf0sLrA",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fce903a2-e4fd-466e-b94d-38d614611ce1"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 288/288 [56:43<00:00, 11.82s/it]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "original_datas = datas\n",
        "# datas = []\n",
        "\n",
        "n_tokens = []\n",
        "\n",
        "datas = []\n",
        "\n",
        "for data in original_datas:\n",
        "    source = data[\"source\"]\n",
        "    text = data[\"text\"]\n",
        "    n_token = len(enc.encode(text))\n",
        "    n_tokens.append(n_token)\n",
        "    if n_token > 8000:\n",
        "        divide_texts = super_divide(text, 6000, 4)\n",
        "        for divide_text in divide_texts:\n",
        "            if len(enc.encode(divide_text)) < 1500:\n",
        "                continue\n",
        "            datas.append({\"source\": source, \"text\": divide_text.strip()})\n",
        "    else:\n",
        "        datas.append({\"source\": source, \"text\": text.strip()})\n"
      ],
      "metadata": {
        "id": "JX1UXXjrCpaZ"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def count_quotes(text):\n",
        "\n",
        "    quote_chars = '''「」\"“”'''\n",
        "    count = 0\n",
        "    for char in text:\n",
        "        if char in quote_chars:\n",
        "            count += 1\n",
        "    return count"
      ],
      "metadata": {
        "id": "Tjq9GTNssyfc"
      },
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def count_paragraphs(text):\n",
        "    lines = text.split(\"\\n\")\n",
        "    count = 0\n",
        "    for line in lines:\n",
        "        if line.strip() != \"\":\n",
        "            count += 1\n",
        "    return count"
      ],
      "metadata": {
        "id": "-yIyDBqvtQTy"
      },
      "execution_count": 55,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm import tqdm\n",
        "n_paragraphs = []\n",
        "n_quotes = []\n",
        "\n",
        "for data in tqdm(datas):\n",
        "    n_paragraphs.append(count_paragraphs(data[\"text\"]))\n",
        "    n_quotes.append(count_quotes(data[\"text\"]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BdNMMR--saEL",
        "outputId": "80e8bd41-6863-41fa-e3fc-a401531b0bab"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 5198/5198 [00:03<00:00, 1601.29it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "对于一个字符串line，检查里面的'“'或者'\"'或者'”'出现位置。并把所有位置合并为一个pos的list，如果pos的长度为偶数，则把奇数位置替换为「把偶数位置替换为」，请用python为我实现"
      ],
      "metadata": {
        "id": "sAFytaQowe1n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def change_quotes(text):\n",
        "    lines = text.split(\"\\n\")\n",
        "\n",
        "    quotes = ['“', '”', '\"']\n",
        "\n",
        "    new_text = \"\"\n",
        "    for line in lines:\n",
        "        if line.strip() == \"\":\n",
        "            continue\n",
        "        positions = []\n",
        "        for i, char in enumerate(line):\n",
        "            if char in quotes:\n",
        "                positions.append(i)\n",
        "        if len(positions) % 2 == 0:\n",
        "            for i in range(len(positions)):\n",
        "                if i % 2 != 0:\n",
        "                    line = line[:positions[i]] + '」' + line[positions[i]+1:]\n",
        "                else:\n",
        "                    line = line[:positions[i]] + '「' + line[positions[i]+1:]\n",
        "        new_text += line + \"\\n\"\n",
        "    return new_text\n"
      ],
      "metadata": {
        "id": "BZIjnTDVwZL4"
      },
      "execution_count": 57,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompt_templates = \"\"\"Proceed with the paragraph, maintaining the same style of language.\n",
        "Keep writing the paragraph, using a similar linguistic style.\n",
        "Continue the paragraph, ensuring the language style remains consistent.\n",
        "Go on with the paragraph in a similar style of language.\n",
        "Carry on writing the paragraph, sticking to the current language style.\n",
        "Extend the paragraph, keeping the language tone and style similar.\n",
        "Resume the paragraph, mirroring the existing language style.\n",
        "Continue crafting the paragraph, maintaining the established language style.\n",
        "Keep extending the paragraph, following the same style of language.\n",
        "Proceed with extending the paragraph, preserving the linguistic style.\n",
        "Please compose a paragraph maintaining a uniform style throughout.\n",
        "Write a paragraph ensuring that the style is consistent from start to finish.\n",
        "Your assignment is to craft a paragraph with a consistent stylistic approach.\n",
        "Create a paragraph where the style remains constant throughout.\n",
        "Construct a paragraph, keeping the writing style uniform in the entire passage.\n",
        "Develop a paragraph that demonstrates consistency in style.\n",
        "Focus on writing a paragraph with a steady and unvarying style.\n",
        "Your challenge is to author a paragraph that maintains a consistent style.\n",
        "Compose a paragraph, ensuring stylistic consistency throughout.\n",
        "Produce a paragraph in which the style does not vary from beginning to end.\n",
        "Your task is to write paragraph in a consistent style\n",
        "对于下面的Paragraph，使用一致性的风格进行续写\n",
        "请继续下面的段落，并保持原有的写作风格一致。\n",
        "在接下来的写作中，延续这一段落的风格，并进行扩展。\n",
        "针对以下段落，以相同的风格进行补充写作。\n",
        "续写下面的段落，并确保风格与原文保持一致。\n",
        "请以与下面段落相同的风格，进行创作性的续写。\n",
        "保持原段落的写作风格，对其进行扩展和深化。\n",
        "延续以下段落的风格，进行相应的续写工作。\n",
        "请在下面的段落后续写，确保风格的连贯性。\n",
        "继续这一段落的故事，注意保持原有风格的一致性。\n",
        "针对下文，以一种风格上连贯的方式进行创作性扩写。\"\"\"\n",
        "\n",
        "prompt_templates = prompt_templates.split(\"\\n\")\n",
        "\n",
        "for i in range(len(prompt_templates)):\n",
        "    prompt_templates[i] = prompt_templates[i].strip()"
      ],
      "metadata": {
        "id": "KRIPcXfDypSY"
      },
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def segment_paragraph(text):\n",
        "    raw_lines = text.split(\"\\n\")\n",
        "    lines = []\n",
        "\n",
        "    all_token_count = 0\n",
        "\n",
        "    for line in raw_lines:\n",
        "        line = line.strip()\n",
        "        if line == \"\":\n",
        "            continue\n",
        "        n_token = len(enc.encode(line))\n",
        "        all_token_count += n_token\n",
        "\n",
        "        if n_token > 600:\n",
        "            segments = super_divide(line)\n",
        "            lines.extend(segments)\n",
        "        else:\n",
        "            lines.append(line)\n",
        "\n",
        "    seg_token = all_token_count // 3\n",
        "\n",
        "    if seg_token > 1500:\n",
        "        seg_token = 1500\n",
        "\n",
        "    current_seg_token = 0\n",
        "    current_seg = \"\"\n",
        "\n",
        "    segments = []\n",
        "    for line in lines:\n",
        "        line = line.strip('\"#-。，, \\n')\n",
        "        if line.strip() == \"\":\n",
        "            continue\n",
        "        n_token = len(enc.encode(line))\n",
        "        if current_seg_token + n_token > seg_token:\n",
        "            if current_seg_token > 0:\n",
        "                segments.append(current_seg.strip())\n",
        "                current_seg = line\n",
        "                current_seg_token = n_token\n",
        "        else:\n",
        "            current_seg += \"\\n\" + line\n",
        "            current_seg_token += n_token\n",
        "\n",
        "    if current_seg_token > 0:\n",
        "        segments.append(current_seg.strip())\n",
        "\n",
        "    return segments"
      ],
      "metadata": {
        "id": "qJ8eQujTz44g"
      },
      "execution_count": 59,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def find_first_sep( text ):\n",
        "    first_sep = 9999\n",
        "    first_sep_second_choice = 9999\n",
        "\n",
        "    find_first_ket = text[:50].find(\"」\",0)\n",
        "\n",
        "    if find_first_ket > 0:\n",
        "        return find_first_ket\n",
        "\n",
        "    if len(text) <= 50:\n",
        "        return len(text)\n",
        "    for i, c in enumerate(text):\n",
        "        if i > 50:\n",
        "            break\n",
        "        if c in ['。','\\n']:\n",
        "            first_sep = min(first_sep, i)\n",
        "        elif c in ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                 '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                 '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                 '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "               '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                 '·', '…']:\n",
        "            first_sep_second_choice = min(first_sep_second_choice, i)\n",
        "\n",
        "    for i in range(5):\n",
        "        if len(text) > first_sep and text[first_sep+1] in ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                    '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                    '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                    '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "                '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                    '·', '…','「','」']:\n",
        "            first_sep += 1\n",
        "\n",
        "    if first_sep <= 50:\n",
        "        return first_sep\n",
        "    elif first_sep_second_choice <= 50:\n",
        "        return first_sep_second_choice\n",
        "    else:\n",
        "        return 50"
      ],
      "metadata": {
        "id": "huG54_dK-NXE"
      },
      "execution_count": 60,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "count = 0\n",
        "\n",
        "import random\n",
        "\n",
        "save_datas = []\n",
        "\n",
        "for i in tqdm(range(len(datas))):\n",
        "\n",
        "    source = datas[i][\"source\"]\n",
        "    # source_pool.add(source)\n",
        "\n",
        "    text = datas[i][\"text\"]\n",
        "    text = change_quotes( text )\n",
        "    segments = segment_paragraph( text )\n",
        "\n",
        "    if len(segments) < 3:\n",
        "        continue\n",
        "\n",
        "    context = random.choice(prompt_templates)\n",
        "\n",
        "    context += \"\\nParagraph:\\n###\\n\" + segments[0] + \"\\n###\\n\"\n",
        "\n",
        "\n",
        "\n",
        "    first_sep = find_first_sep(segments[1])\n",
        "\n",
        "    context += segments[1][:first_sep+1]\n",
        "\n",
        "    target = segments[1][first_sep+1:]\n",
        "\n",
        "    for j in range(2, len(segments)):\n",
        "        target += \"\\n\\n\" + segments[j]\n",
        "\n",
        "    save_data = {\n",
        "        \"context\": context,\n",
        "        \"target\": target\n",
        "    }\n",
        "\n",
        "    save_datas.append(save_data)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gxHPWRf_wQOd",
        "outputId": "ec765a24-5179-48d7-aa0b-a69566531663"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 5198/5198 [02:45<00:00, 31.41it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(len(save_datas))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fPEpg9UvETx4",
        "outputId": "efd16461-55ae-4e19-df69-ad6da74a5ff8"
      },
      "execution_count": 62,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1621\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "save_name = \"/content/erotics.jsonl\"\n",
        "\n",
        "with open(save_name, 'w', encoding='utf8') as f:\n",
        "    for data in save_datas:\n",
        "        data['target'] = data['target'].strip()\n",
        "        json.dump(data, f, ensure_ascii=False)\n",
        "        f.write('\\n')"
      ],
      "metadata": {
        "id": "UJDMmlQg0ulC"
      },
      "execution_count": 67,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(len(save_datas))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1XqAn14QFb-d",
        "outputId": "23ed73fc-25e8-4c4f-b4d6-94186a769ff8"
      },
      "execution_count": 68,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1621\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "UglKKm1Wcmdr"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}